# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import os
import numpy as np
import pandas as pd

from absl import flags, app
from absl.flags import FLAGS


class BalanceInfo(object):
  def __init__(self,
               time,
               quote0_name,
               quote0,
               quote1_name,
               quote1,
               quote0_rate,
               quote1_rate,
               imbalance,
               estimatedbase,
               pnl_quote_name,
               okdk):
    self.time = time
    self.quote0_name = quote0_name
    self.quote0 = quote0
    self.quote1_name = quote1_name
    self.quote1 = quote1
    self.quote0_rate = quote0_rate
    self.quote1_rate = quote1_rate
    self.imbalance = imbalance
    self.estimatedbase = estimatedbase
    self.pnl_quote_name = pnl_quote_name
    self.okdk = okdk

  def total_value(self):
    quote0 = self.quote0
    quote1 = self.quote1
    if self.okdk != 0:
      if 'USD' in self.quote0_name.upper():
        quote0 += self.okdk * FLAGS.okdk_price
      elif 'USD' in self.quote1_name.upper():
        quote1 += self.okdk * FLAGS.okdk_price
      else:
        assert False, 'Okdk %f, %s, %s' % \
            (self.okdk, self.quote0_name, self.quote1_name)

    return quote0 * self.quote0_rate + \
        (quote1 + self.imbalance) * self.quote1_rate

  def output(self):
    adjusted_quote1 = self.quote1 + self.imbalance
    estimated_total = self.total_value()
    print(
        '''%s
%s: %f
%s: %f
Imbalance: %f
Adjusted%s: %f
Okdk: %f
EstimatedTotal%s: %f
EstimatedBase: %f''' \
    % (self.time, self.quote0_name, self.quote0, self.quote1_name, self.quote1,
       self.imbalance, self.quote1_name, adjusted_quote1, self.okdk,
       self.pnl_quote_name, estimated_total, self.estimatedbase))


def find_balances(log_dir,
                  archive_log_dir,
                  from_dt,
                  to_dt,
                  quote0_rate,
                  quote1_rate,
                  pnl_quote_name):
  from_balance = find_a_balance(log_dir,
                                archive_log_dir,
                                from_dt,
                                quote0_rate,
                                quote1_rate,
                                pnl_quote_name)
  to_balance = find_a_balance(log_dir,
                              archive_log_dir,
                              to_dt,
                              quote0_rate,
                              quote1_rate,
                              pnl_quote_name)
  return from_balance, to_balance


def find_a_balance(log_dir, archive_log_dir, cutoff, quote0_rate, quote1_rate, pnl_quote_name):
  logs = [os.path.join(log_dir, log) for log in os.listdir(log_dir) if not log.startswith('.')]
  assert len(logs) == 1, logs
  archive_logs = []
  if os.path.isdir(archive_log_dir):
    for a_log_dir in sorted(os.listdir(archive_log_dir), reverse=True):
      if a_log_dir.startswith(os.path.basename(log_dir)):
        a_log_dir = os.path.join(archive_log_dir, a_log_dir)
        if os.path.isdir(a_log_dir):
          archive_logs += [
              os.path.join(a_log_dir, log)
              for log in os.listdir(a_log_dir)
              if not log.startswith('.') and not log.endswith('.zip')
          ]
  logs = logs + archive_logs
  for log in logs:
    if os.path.isdir(log):
      continue
    balance_info = find_balance_in_log(cutoff, log, quote0_rate, quote1_rate, pnl_quote_name)
    if balance_info is not None:
      return balance_info


def find_balance_in_log(cutoff, log_path, quote0_rate, quote1_rate, pnl_quote_name):
  has_eariler = False
  found = False
  time = ''
  with open(log_path, 'r', encoding='utf8') as log:
    for line in log.readlines():
      line = line.strip()
      if found:
        if ':' in line:
          data = [
              (item.split(':')[0].strip(), float(item.split(':')[1])) for item in line.split(',')
          ]
          data_map = {d[0]: d[1] for d in data}
          okdk = data_map.get('OKDK', 0)
          if data[1][0] == 'Imbalance':
            balance_info = BalanceInfo(time,
                                       data[0][0],
                                       0.0,
                                       data[0][0],
                                       data[0][1],
                                       quote0_rate,
                                       quote1_rate,
                                       data[1][1],
                                       data[3][1],
                                       pnl_quote_name,
                                       okdk)
          else:
            balance_info = BalanceInfo(time,
                                       data[0][0],
                                       data[0][1],
                                       data[1][0],
                                       data[1][1],
                                       quote0_rate,
                                       quote1_rate,
                                       data[2][1],
                                       data[4][1],
                                       pnl_quote_name,
                                       okdk)
          return balance_info
      if line.startswith('*'):
        fmt = '%Y-%m-%d %H:%M:%S.%f' if '.' in line else '%Y-%m-%d %H:%M:%S'
        a_datetime = datetime.datetime.strptime(line.split(' ', 1)[1], fmt)
        if a_datetime < cutoff:
          has_eariler = True
        else:
          if has_eariler:
            found = True
            time = line
          else:
            return None


def filter_fills_df_by_strategy_name(df, strategy_names):
  mask = df['strategy_name'].isin(strategy_names)
  return df[mask]


def filter_fills_df_by_fills(df):
  fill_mask = df['action'].str.contains('filled') | \
      df['action'].str.contains('Filled')
  return df[fill_mask]


def filter_fills_df_by_period(df, from_dt, to_dt):
  within_mask = (df['timestamp'] > from_dt.timestamp() * 10 ** 9) & \
      (df['timestamp'] <= to_dt.timestamp() * 10 ** 9)
  return df[within_mask]


def filter_invalid_rows(df):
  if df['fill_price'].dtype == np.object_:
    none_mask = df['fill_price'] == 'None'
    df = df[~none_mask]
    df.loc[:, 'fill_price'] = df['fill_price'].astype(float)
    df.loc[:, 'fill_qty'] = df['fill_qty'].astype(float)
  return df


def update_price_for_futures(df):
  is_future_mask = df['symbol'].str.contains('\.')
  is_btc_future_mask = df['symbol'].str.contains('BTC') & is_future_mask
  df.loc[is_future_mask, 'fill_price'] = 10.0
  df.loc[is_btc_future_mask, 'fill_price'] = 100.0
  return df


def find_volume_df(working_dir, from_dt, to_dt, strategy_name=None):
  df = pd.read_csv(os.path.join(working_dir, 'fills.csv'))
  if strategy_name:
    strategies = strategy_name.split(',')
    df = filter_fills_df_by_strategy_name(df, strategies)
  df = filter_fills_df_by_fills(df)
  df = filter_fills_df_by_period(df, from_dt, to_dt)
  df = filter_invalid_rows(df)
  df = update_price_for_futures(df)
  df = df.reset_index(drop=True)
  return df


def find_volumes(df, quote0, quote1):
  if quote0 == quote1:
    mask = df['symbol'].str.contains(quote0)
    return {quote0: np.sum(df.loc[mask, 'fill_price'] * df.loc[mask, 'fill_qty'])}
  else:
    mask0 = df['symbol'].str.endswith(quote0)
    mask1 = df['symbol'].str.endswith(quote1)
    assert np.sum(mask0) + np.sum(mask1) == df.shape[0]
    return {
        quote0: np.sum(df.loc[mask0, 'fill_price'] * df.loc[mask0, 'fill_qty']),
        quote1: np.sum(df.loc[mask1, 'fill_price'] * df.loc[mask1, 'fill_qty'])
    }


def find_volume_breakdown(df):
  df['volume'] = df['fill_price'] * df['fill_qty']
  return df.groupby('symbol')['volume'].sum()


def main(args):
  if FLAGS.trading_date:
    trading_date = datetime.datetime.strptime(FLAGS.trading_date, '%Y%m%d')
  else:
    trading_date = datetime.datetime.today().replace(hour=0, minute=0, second=0, microsecond=0)
  hours, minutes = FLAGS.cutoff_hour_min.split(':')
  cutoff_timedelta = datetime.timedelta(hours=int(hours), minutes=int(minutes))
  to_dt = trading_date + cutoff_timedelta
  from_dt = to_dt - datetime.timedelta(hours=FLAGS.duration_hours)

  working_dirs = [os.path.expanduser(path) for path in FLAGS.working_dir.split(',')]
  log_dir = os.path.join(working_dirs[0], FLAGS.current_log_dir)
  archive_log_dir = os.path.join(working_dirs[0], FLAGS.archive_log_dir)
  from_bal, to_bal = find_balances(log_dir, archive_log_dir, from_dt, to_dt, FLAGS.quote0_rate,
                                   FLAGS.quote1_rate, FLAGS.pnl_quote)

  volumes = dict()
  volumes_breakdown = None
  for volume_dir in working_dirs:
    volume_df = find_volume_df(volume_dir, from_dt, to_dt, FLAGS.strategy_name)
    a_volumes = find_volumes(volume_df, from_bal.quote0_name, from_bal.quote1_name)
    for ticker, volume in a_volumes.items():
      volumes[ticker] = volumes.get(ticker, 0.0) + volume
    a_volumes_breakdown = find_volume_breakdown(volume_df)
    volumes_breakdown = a_volumes_breakdown if volumes_breakdown is None \
        else volumes_breakdown.append(a_volumes_breakdown)

  from_bal.output()
  print('')
  to_bal.output()
  print('\nPnl:\n%f' % (to_bal.total_value() - from_bal.total_value()))
  for key, val in volumes.items():
    print('Volume %s:\n%f' % (key, val))
  print('\nVolume Breakdown:')
  for ticker, volume in volumes_breakdown.iteritems():
    print('%s: %f' % (ticker, volume))


if __name__ == '__main__':
  flags.DEFINE_string('working_dir', '.', '')
  flags.DEFINE_string('current_log_dir', 'log', '')
  flags.DEFINE_string('archive_log_dir', 'archive', '')
  flags.DEFINE_string('trading_date', '', '')
  flags.DEFINE_string('cutoff_hour_min', '00:00', 'UTC time')
  flags.DEFINE_float('duration_hours', 24.0, '')
  flags.DEFINE_float('quote0_rate', 1, '')
  flags.DEFINE_float('quote1_rate', 1, '')
  flags.DEFINE_string('pnl_quote', 'USD', '')
  flags.DEFINE_string('strategy_name', '', '')
  flags.DEFINE_float('okdk_price', 0.4, '')

  app.run(main)
